package heteroset;

import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.Set;

public class HeteroHashSet<E> implements HeteroSet<E> {
	private final HashMap<Class<? extends E>, Set<?>> map = new HashMap<>();
	private int size = 0;

	@Override
	public <T extends E> boolean add(Class<T> type, T elem) {
		if (type == null)
			throw new NullPointerException();
		@SuppressWarnings("unchecked")
		Set<T> s = (Set<T>) map.get(type);
		if (s == null) {
			s = new HashSet<>();
			map.put(type, s);
		}
		boolean result = s.add(elem);
		
		if (result)
			++size;
		return result;
	}
	
	@Override
	public <T extends E> boolean remove(Class<T> type, T elem) {
		if (type == null)
			throw new NullPointerException();
		@SuppressWarnings("unchecked")
		Set<? extends E> s = (Set<? extends E>) map.get(type);
		
		boolean result = s!=null && s.remove(elem);
		
		if (result)
			--size;
		if (s.isEmpty())
			map.remove(type);
		
		return result;
		
	}

	@Override
	public <T extends E> boolean contains(Class<T> type, T elem) {
		if (type == null)
			throw new NullPointerException();

		@SuppressWarnings("unchecked")
		Set<? extends E> s = (Set<? extends E>) map.get(type);
		return s != null && s.contains(elem);
	}

	@Override
	public boolean isEmpty() {
		return map.isEmpty();
	}

	@Override
	public int size() {
		return size;
	}

	@Override
	public <T extends E> Iterator<T> iterator(Class<T> type) {
		if (type == null)
			throw new NullPointerException();
		@SuppressWarnings("unchecked")
		Set<T> s = (Set<T>) map.get(type);
		if (s == null)
			s = new HashSet<>();
		return s.iterator();
	}
	
	@SuppressWarnings("unchecked")
	@Override
	public Iterator<E> iterator() {
		LinkedList<E> l = new LinkedList<>();
		
		for (Set<?> set : map.values())
			l.addAll((Collection<E>)set);
		
		return l.iterator();
	}

}
